# -*- coding: utf-8 -*-


import os
import time
import pandas as pd
from leaky_lstm_model import cnn_lstm_leak_model
from lstm_model import make_tokenizer, pad_sequence, make_embedding_matrix


def average_fusion(results, model_cnt=5):
    return sum(results) / model_cnt


if __name__ == '__main__':
    appMaxLen = 174  # 1000
    appEmbedDim = 50
    appTknz = make_tokenizer(os.path.join("data", "deviceid_packages_cooked.txt"),
                             os.path.join("cache", "app.tokenizer"))
    appEmbeddings = make_embedding_matrix(appTknz, appEmbedDim, os.path.join("cache", "app_cbow_50d.vec"),
                                          os.path.join("cache", "app"))

    brandMaxLen = 10
    brandEmbedDim = 5
    brandTknz = make_tokenizer(os.path.join("data", "deviceid_brand_cooked.txt"),
                               os.path.join("cache", "brand.tokenizer"))
    brandEmbeddings = make_embedding_matrix(brandTknz, brandEmbedDim, os.path.join("cache", "brand_cbow_5d.vec"),
                                            os.path.join("cache", "brand"))

    df = pd.read_csv(os.path.join("data", "app_brand.test"), sep=",", names=["person", "app", "brand", "label"])
    print df.shape
    ids = df.pop("person")

    leak_df = pd.read_csv(os.path.join("data", "app_gender.test"), sep=",")
    # print leak_df.head()
    df_mean = pd.read_pickle(os.path.join("cache", "age_mean.pkl"))
    # print df_mean
    df_std = pd.read_pickle(os.path.join("cache", "age_std.pkl"))
    # print df_std
    leak_names = leak_df.columns.values
    leak_df = (leak_df - df_mean) / df_std

    leakDim = leak_df.shape[1]
    print leak_df.shape
    # print leak_df.head()

    df = pd.concat([df, leak_df], axis=1)
    print df.shape
    # """
    test_input = [pad_sequence(df["app"], appTknz, appMaxLen), pad_sequence(df["brand"], brandTknz, brandMaxLen),
                  df[leak_names]]

    modelCnt = 5
    probs = []
    for i in xrange(modelCnt):
        languageModel = cnn_lstm_leak_model(
            appTknz, appEmbedDim, appEmbeddings, appMaxLen,
            brandTknz, brandEmbedDim, brandEmbeddings, brandMaxLen,
            leakDim,
            num_class=22
        )

        print "Load model {} ...".format(i)
        languageModel.load_weights((os.path.join("model", "cnn_lstm_leak_{}.h5".format(i))))

        probs.append(languageModel.predict(test_input, batch_size=256, verbose=1))

    writer = open(os.path.join("submission", "result_{}.csv".format(time.strftime('%Y%m%d', time.localtime()))), "w")
    header = []
    for i in xrange(2):
        for j in xrange(11):
            header.append("{}-{}".format(i + 1, j))
    writer.write("DeviceID, {}\n".format(", ".join(header)))

    cnt = 0
    for prob in average_fusion(probs):
        # print sum(prob)
        writer.write("{}, {}\n".format(ids[cnt], ", ".join(map(str, prob))))
        cnt += 1
    print "Done."
    # """
